//-----------------------------------------------------------------------------
//
// @file	LevelSet.cpp
// @brief	LevelSet file
// @author	Vinod Melapudi
// @date	13Aug2011
//
//-----------------------------------------------------------------------------

#include "LevelSet.h"
#include "MemPool.h"
#include "MarchingCubes.h"
#include "Scene.h"

#include "GL/gl.h"
#include "GL/glu.h"

// ------------------------------------------------------------------------------------

#define GET_NEXT_NODE_HELPER( _node )	(_node && _node->next && _node->next->idx == _node->idx+1) ? _node->next : NULL;

#define GET_INFO_HELPER( _z_node, _target )						\
		if ( _z_node )											\
		{														\
			InfoNode* i = (InfoNode*) _z_node->param;			\
			assert(i);											\
			_target = i->value;									\
		}

// ------------------------------------------------------------------------------------


// Structure used for Rendering kernel function
struct RenderUserData
{
	// Render to mesh params
	MarchingCubes::Input	*input;
	MarchingCubes::Output	*output;
	Mesh					*mesh;

	// Render grid params
	Scene					*scene;
	vector<Vector3>			cell_pos;
	vector<SceneText>		value_texts;
	vector<SceneText>		highlighted_texts;
};


// ------------------------------------------------------------------------------------

LevelSet::LevelSet()
: m_octree(NULL)
, m_mem_info_pool(NULL)
, m_cell_size(1.f,1.f,1.f)
, m_highlighted_cell(0,0,0)
{
}

LevelSet::~LevelSet()
{
	Reset();
}

// Reset level set
void LevelSet::Reset()
{
	if ( m_mem_info_pool )
		delete m_mem_info_pool;

	delete m_octree;
	m_octree = NULL;

	m_mem_info_pool = NULL;

	m_cell_size.set(1.f,1.f,1.f);
}

// Allocate a info node, 
LevelSet::InfoNode*		LevelSet::AllocateInfoNode( float value )
{
	assert( m_mem_info_pool );
	InfoNode* n = (InfoNode*) m_mem_info_pool->Allocate( sizeof(InfoNode) );
	n->value = value;
	n->new_value = value;
	return n;
}


// Get info node of the cell
LevelSet::InfoNode* LevelSet::GetInfoNode( const uLocation& l )
{
	if ( !m_octree )
		return NULL;

	InfoNode** info_ptr = m_octree->Get(l);
	if ( info_ptr )
		return *info_ptr;

	return NULL;
}




// Set value at grid point
void LevelSet::SetValue( const uLocation& l, float value )
{
	if ( !m_octree )
		return;

	// Check if info node already exists at that place
	InfoNode* info = GetInfoNode( l );
	if ( info )
	{
		info->value = value;
		info->new_value = value;
		return;
	}

	// Create new node
	info = AllocateInfoNode( value );
	m_octree->Set( l, info );
}


// Get value at grid point (to support float and double)
template<class T>
bool LevelSet::GetValue( const uLocation& l, T &value )
{
	InfoNode* info = GetInfoNode( l );
	if ( info )
	{
		value = info->value;
		return true;
	}
	return false;
}




// Get neighbour values
bool LevelSet::GetNeighbourValues(	const uLocation& l, float& node_value, 
									Vector3* pos_values, Vector3* neg_values, Vector4* diag_values )
{
	// Get node values
	float val;
	bool valid = GetValue( l, val );
	if ( valid )		
		node_value = val;		
	else	
		return false;

	// Get positive neighbour values
	if ( pos_values )
	{
		Vector3 pos;
		if (	GetValue( uLocation(l.x+1, l.y, l.z), pos.x ) &&
				GetValue( uLocation(l.x, l.y+1, l.z), pos.y ) &&
				GetValue( uLocation(l.x, l.y, l.z+1), pos.z ) )
			*pos_values = pos;		
		else	
			return false;
	}

	// Get negative neighbour values
	if ( neg_values )
	{
		Vector3 neg;
		if (	GetValue( uLocation(l.x-1, l.y, l.z), neg.x ) &&
				GetValue( uLocation(l.x, l.y-1, l.z), neg.y ) &&
				GetValue( uLocation(l.x, l.y, l.z-1), neg.z ) )
			*neg_values = neg;
		else	
			return false;
	}

	// Get diagonal values
	if ( diag_values )
	{
		// Get neighbours of x+1,y+1,z+1
		float d_val;
		Vector3 d_neg_values;
		if ( GetNeighbourValues( uLocation(l.x+1, l.y+1, l.z+1), d_val, NULL, &d_neg_values, NULL ) )
			diag_values->set( d_neg_values.x, d_neg_values.y, d_neg_values.z, d_val );
		else
			return false;
	}

	return true;
}


// Get neighbour values diff
bool LevelSet::GetNeighbourDiff( const uLocation& l, Vector3& pos_diff, Vector3& neg_diff )
{
	Vector3 pos_values;
	Vector3 neg_values;
	float value;

	if ( GetNeighbourValues( l, value, &pos_values, &neg_values ) )
	{
		pos_diff.set ( pos_values.x - value, pos_values.y - value, pos_values.z - value );
		neg_diff.set ( value - neg_values.x, value - neg_values.y, value - neg_values.z );
		return true;
	}

	return false;
}


// Grow level set
void LevelSet::Grow( float factor )
{
	if ( factor == 0.f )
		return;

	ApplyFn( &LevelSet::GrowCellFn, (void*)&factor );
	ApplyFn( &LevelSet::ReplaceValueCellFn );
	Normalize();
}

// Apply a vector field to level set
void LevelSet::ApplyVectorField( VectorFieldFn fn )
{
	assert( fn );
	ApplyFn( &LevelSet::ApplyVectorFieldCellFn, (void*)fn );
	ApplyFn( &LevelSet::ReplaceValueCellFn );
	Normalize();
}



// Normalize level set
void LevelSet::Normalize()
{
//	ApplyFn( &LevelSet::NormalizeCellFn );
//	ApplyFn( &LevelSet::ReplaceValueCellFn );
}
	
// Set level set to sphere
void LevelSet::SetSphere( float radius, float cell_size )
{
	int extra_cells_on_1_side = 3;
	int num_cells_to_fill_on_1_side = radius / cell_size + extra_cells_on_1_side;

	int num_cells_per_axis = NearestPowerOf2( num_cells_to_fill_on_1_side * 2.f );

	// Allocated mem
	Reset();
	m_octree = new Octree<InfoNode*>( num_cells_per_axis );
	m_mem_info_pool = new MemPool( sizeof(InfoNode), 0, num_cells_per_axis*num_cells_per_axis ); 


	// Set cell size
	m_cell_size.set( cell_size, cell_size, cell_size );

	float center = (num_cells_per_axis / 2) * cell_size;
	float normalize_factor = 1 / sqrt( center*center*3 );


	// Start cell to fill info
//	int start_cell_to_fill = num_cells_per_axis / 2 - num_cells_to_fill_on_1_side;
//	int end_cell_to_fill = num_cells_per_axis / 2 + num_cells_to_fill_on_1_side;

	int start_cell_to_fill = 0;
	int end_cell_to_fill = num_cells_per_axis;

	// Fill up grid
	for( int x=start_cell_to_fill; x < end_cell_to_fill; x++ )
	{
		// Distance from center
		float xdist = x*cell_size - center;
		float xdist2 = xdist*xdist;

		for( int y=start_cell_to_fill; y < end_cell_to_fill; y++ )
		{
			// Distance from center
			float ydist = y*cell_size - center;
			float ydist2 = ydist*ydist;

			for( int z=start_cell_to_fill; z < end_cell_to_fill; z++ )
			{
				// Distance from center
				float zdist = z*cell_size - center;
				float dist2 = xdist2 + ydist2 + zdist*zdist;

				float value = sqrt(dist2) - radius;
				value *= normalize_factor;			// Normalize
				SetValue( uLocation(x, y, z), value );
			}
		}
	}
}

// Render level set to mesh using marching cubes algo
void LevelSet::RenderToMesh( Mesh& mesh )
{
	mesh.Reset();

	// Initialize marching cube algo input/output;
	MarchingCubes::Input input;
	MarchingCubes::Output output;
	input.isolevel = 0.f;
	input.grid_pts[0].set( 0,				0,				0 );
	input.grid_pts[1].set( m_cell_size.x,	0,				0 );
	input.grid_pts[2].set( m_cell_size.x,	0,				m_cell_size.z );
	input.grid_pts[3].set( 0,				0,				m_cell_size.z );
	input.grid_pts[4].set( 0,				m_cell_size.y,	0 );
	input.grid_pts[5].set( m_cell_size.x,	m_cell_size.y,	0 );
	input.grid_pts[6].set( m_cell_size.x,	m_cell_size.y,	m_cell_size.z );
	input.grid_pts[7].set( 0,				m_cell_size.y,	m_cell_size.z );

	// Apply kernel fn to grid
	RenderUserData rud;
	rud.input = &input;
	rud.output = &output;
	rud.mesh = &mesh;
	ApplyFn( &LevelSet::RenderToMeshCellFn, (void*)(&rud) );
}


// Render cell info (cell lines and values)
void LevelSet::RenderGrid( Renderer* renderer, Scene* scene )
{
	// Apply kernel fn to grid
	RenderUserData rud;
	rud.scene = scene;
	ApplyFn( &LevelSet::RenderGridCellFn, (void*)(&rud) );

	// Render grid lines
	glColor3f( 0.f, 0.2f, 0.2f );
	glBegin( GL_LINES );
	for( int i=0; i < rud.cell_pos.size(); i++ )
	{
		Vector3& pos = rud.cell_pos[i];
		glVertex3f( pos.x,					pos.y,					pos.z );
		glVertex3f( pos.x+m_cell_size.x,	pos.y,					pos.z );
		glVertex3f( pos.x,					pos.y,					pos.z );
		glVertex3f( pos.x,					pos.y+m_cell_size.y,	pos.z );
		glVertex3f( pos.x,					pos.y,					pos.z );
		glVertex3f( pos.x,					pos.y,					pos.z+m_cell_size.z );
	}
	glEnd();

	// Render the texts
	glColor3f( 0.f, 0.2f, 0.2f );
	scene->RenderText( renderer, rud.value_texts );
	glColor3f( 0.f, 0.5f, 0.3f );
	scene->RenderText( renderer, rud.highlighted_texts );
}


// Apply function(Kernel) to each grid node
void LevelSet::ApplyFn( CellOperationFn fn, void* user_data )
{
	CellOpParam c;
	c.level_set = this;
	c.user_data = user_data;

	// Run through grid and apply cell operation with user data
	for( Octree<InfoNode*>::iterator iter = m_octree->Begin(); iter < m_octree->End(); iter++ )
	{
		c.location = iter.GetLocation();
		c.info_node = **iter;
		(*fn)( c );
	}
}

// Helper fn to calc gradient length based on given factor
bool LevelSet::CalcGradientLength( const uLocation& l, float factor, float &gradient_length )
{
	// Get diffs with neighbour values
	Vector3 pos_diff, neg_diff;
	if ( GetNeighbourDiff( l, pos_diff, neg_diff ) == false )
		return false;

	// Get derivative
	pos_diff.x /= m_cell_size.x;
	pos_diff.y /= m_cell_size.y;
	pos_diff.z /= m_cell_size.z;
	neg_diff.x /= m_cell_size.x;
	neg_diff.y /= m_cell_size.y;
	neg_diff.z /= m_cell_size.z;

	// Calc gradient length
	float gradient_length_sqr = 0.f;
/*
	if ( factor < 0 )
		gradient_length_sqr =	max( pow( (float)max( neg_diff.x, 0 ), 2.f ), pow( (float)min( pos_diff.x, 0 ), 2.f ) ) +
								max( pow( (float)max( neg_diff.y, 0 ), 2.f ), pow( (float)min( pos_diff.y, 0 ), 2.f ) ) +
								max( pow( (float)max( neg_diff.z, 0 ), 2.f ), pow( (float)min( pos_diff.z, 0 ), 2.f ) );
	else if ( factor > 0 )
		gradient_length_sqr =	max( pow( (float)min( neg_diff.x, 0 ), 2.f ), pow( (float)max( pos_diff.x, 0 ), 2.f ) ) +
								max( pow( (float)min( neg_diff.y, 0 ), 2.f ), pow( (float)max( pos_diff.y, 0 ), 2.f ) ) +
								max( pow( (float)min( neg_diff.z, 0 ), 2.f ), pow( (float)max( pos_diff.z, 0 ), 2.f ) );
	gradient_length = sqrt( gradient_length_sqr );
*/

	Vector3 pos_factor = factor * pos_diff;
	Vector3 neg_factor = factor * neg_diff;

	Vector3 gradient(0,0,0);
	for( int i=0; i < 3; i++ )
	{
		if ( neg_factor[i] <= 0 && pos_factor[i] <= 0 )			gradient[i] = pos_diff[i];
		else if ( neg_factor[i] >= 0 && pos_factor[i] >= 0 )	gradient[i] = neg_diff[i];
		else if ( neg_factor[i] <= 0 && pos_factor[i] >= 0 )	gradient[i] = 0.f;
		else if ( neg_factor[i] >= 0 && pos_factor[i] <= 0 )
		{
			if ( fabs(pos_factor[i]) >= fabs(neg_factor[i]) )
				gradient[i] = pos_diff[i];
			else
				gradient[i] = neg_diff[i];
		}
	}
	
	
	gradient_length = gradient.norm();

	return true;
}

// Grow cell by factor kernel
void LevelSet::GrowCellFn( const CellOpParam& param )
{
	float factor = *((float*)param.user_data);

	// Get gradient length based on factor
	float gradient_length = 0.f;
	if ( param.level_set->CalcGradientLength( param.location, factor, gradient_length ) )
	{
		// Get new value of level set
		param.info_node->new_value = param.info_node->value - gradient_length * factor;
	}
}

// Apply vector field by factor kernel
void LevelSet::ApplyVectorFieldCellFn( const CellOpParam& param )
{
	VectorFieldFn fn = (VectorFieldFn)(param.user_data);

	// Get vector field value and diff of neighbour values
	Vector3 vf_value = 0.f;
	Vector3 pos_diff, neg_diff;
	if ( (*fn)( param.location, vf_value ) && param.level_set->GetNeighbourDiff( param.location, pos_diff, neg_diff ) )
	{
		// Calc gradient
		Vector3 gradient(0,0,0);

		if ( vf_value.x < 0 )		gradient.x = pos_diff.x;
		else if ( vf_value.x > 0 )	gradient.x = neg_diff.x;

		if ( vf_value.y < 0 )		gradient.y = pos_diff.y;
		else if ( vf_value.y > 0 )	gradient.y = neg_diff.y;

		if ( vf_value.z < 0 )		gradient.z = pos_diff.z;
		else if ( vf_value.z > 0 )	gradient.z = neg_diff.z;


		// Get derivative
		const Vector3& cell_size = param.level_set->GetCellSize();
		gradient.x /= cell_size.x;
		gradient.y /= cell_size.y;
		gradient.z /= cell_size.z;

		// Get new values of level set
		param.info_node->new_value = param.info_node->value - gradient * vf_value;
	}
}

// Normalize kernel fn
void LevelSet::NormalizeCellFn( const CellOpParam& param )
{
	float v = param.info_node->value;
	float sgn = v / sqrt( v*v );

	// Get gradient length using smoothed sign value
	float gradient_length = 0.f;
	if ( param.level_set->CalcGradientLength( param.location, -sgn, gradient_length ) )
	{
		param.info_node->new_value = param.info_node->value + sgn * (1 - gradient_length );
	}
}

// Replace old_value with new value (kernel fn)
void LevelSet::ReplaceValueCellFn( const CellOpParam& param )
{
	param.info_node->value = param.info_node->new_value;
}

// Render each cell (kernel fn)
void LevelSet::RenderToMeshCellFn( const CellOpParam& param )
{
	RenderUserData* rud = (RenderUserData*) param.user_data;
	MarchingCubes::Input &input = *rud->input;
	MarchingCubes::Output &output = *rud->output;
	Mesh& mesh = *rud->mesh;

	LevelSet *level_set = param.level_set;

	uint x = param.location.x;
	uint y = param.location.y;
	uint z = param.location.z;

	// Get grid and neighbour signed dist values
	float value;
	Vector3 pos_values;
	Vector4 diag_values;
	if ( !level_set->GetNeighbourValues( param.location, value, &pos_values, NULL, &diag_values ) )
		return;

	// Prep grid values in order given in MarchingCubes.h
	input.grid_values[0] = value;
	input.grid_values[1] = pos_values.x;
	input.grid_values[2] = diag_values.y;
	input.grid_values[3] = pos_values.z;
	input.grid_values[4] = pos_values.y;
	input.grid_values[5] = diag_values.z;
	input.grid_values[6] = diag_values.w;
	input.grid_values[7] = diag_values.x;

	// Run marching cubes algo to get triangles
	if ( MarchingCubes::Polygonize( input, output ) == false )
		return;

	const Vector3& cell_size = level_set->GetCellSize();
	Vector3 cell_origin( cell_size.x * x, cell_size.y * y, cell_size.z * z );

	// Add vertices to mesh
	int vtx_start = mesh.m_vtx_positions.size();
	for( int i=0; i < output.num_vertices; i++ )
		mesh.m_vtx_positions.push_back( cell_origin + output.vertices[i] );

	// Add polygon indices
	for( int i=0; i < output.num_triangles; i++ )
		for( int j=0; j < 3; j++ )
			mesh.m_face_vtx_indices.push_back( vtx_start + (int)output.triangle_vtx_idices[i*3+j] );
}


// Render grid
void LevelSet::RenderGridCellFn( const CellOpParam& param )
{
	// Render params
	RenderUserData* rud = (RenderUserData*) param.user_data;
	Scene* scene = rud->scene;

	// Grid params
	uint x = param.location.x;
	uint y = param.location.y;
	uint z = param.location.z;

	// left set info
	LevelSet *level_set = param.level_set;
	const Vector3& cell_size = level_set->GetCellSize();
	uLocation hcell = level_set->m_highlighted_cell;

	
	// Store grid point
	Vector3 wpos( x * cell_size.x, y * cell_size.y, z * cell_size.z );
	rud->cell_pos.push_back( wpos );

	// Get screen dimension
	int sw, sh;
	scene->GetScreenDimension( sw, sh );
	Vector2 scr_dim( sw/2.f, sh/2.f );


	// Get screen pos
	Vector2 spos;
	char str[16];
	if ( scene->WorldToScreen( wpos, spos ) )
	{
		sprintf( str, "%.2f", param.info_node->value );
		SceneText st( str, spos.x/scr_dim.x, spos.y/scr_dim.y );
		if ( ( x == hcell.x || x == hcell.x+1 ) && 
			 ( y == hcell.y || y == hcell.y+1 ) && 
			 ( z == hcell.z || z == hcell.z+1 ) )
				rud->highlighted_texts.push_back( st );
		//else
		//		rud->value_texts.push_back( st );
	}

}

// Set the cell to be highlighted 
void LevelSet::MoveHighlightCell( const uLocation& move_by )
{
	uLocation highlighted_cell = m_highlighted_cell + move_by;

	float value;
	Vector3 pos_values;
	Vector4 diag_values;
	if ( GetNeighbourValues( highlighted_cell, value, &pos_values, NULL, &diag_values ) )
		m_highlighted_cell = highlighted_cell;
}